#!/usr/bin/env python3
# @file terminal_server.py
# @author HKUST Aerial Robotics Group
# @brief terminal server for testing
# @version 0.1
# @date 2019-02
# @copyright Copyright (c) 2019

import os
from ament_index_python.packages import get_package_share_directory

# sys
import time
import sys
import rclpy
from rclpy.node import Node
import shutil
import random
from math import *
from functools import reduce
# pygame
import pygame as pg
from pygame.locals import *
from pygame.math import Vector2
# msg
from geometry_msgs.msg import Twist
from sensor_msgs.msg import Joy
from vehicle_msgs.msg import ArenaInfoDynamic
from vehicle_msgs.msg import ArenaInfoStatic
from vehicle_msgs.msg import State
from vehicle_msgs.msg import ControlSignal


# Helper function for compatibility
def print_over_same_line(text):
    terminal_width = shutil.get_terminal_size((80, 20)).columns
    empty_space = max(0, terminal_width - len(text))
    sys.stdout.write('\r' + text + empty_space * ' ')
    sys.stdout.flush()


# Mixin class containing methods for TerminalServerNode
class TerminalServerNodeMethods:
    """Mixin class containing methods for TerminalServerNode"""
    
    def calc_current_steer_acc(self):
        if len(self.state_seq) < 2:
            return 0.0, 0.0
        steer_list = []
        for i in range(len(self.state_seq)-1):
            state1 = self.state_seq[i]
            steer = atan(state1.curvature * 2.85)
            steer_list.append(steer)
        ave_steer = reduce(lambda x, y: x + y, steer_list) / len(steer_list)
        rotated_angle = ave_steer * 1.25 * (360.0/45.0)
        return rotated_angle, self.state_seq[-1].acceleration

    def plot_lanes_on_screen(self):
        lanepts_plot = list(self.lane_pts)
        for i in range(len(lanepts_plot)):
            pg.draw.lines(self.screen, pg.Color('deeppink'), False, lanepts_plot[i])

    def plot_speed_on_screen(self):
        speed = 0.0
        acc = 0.0
        if len(self.state_seq) > 1:
            speed = self.state_seq[-1].velocity * 3.6
            acc = self.state_seq[-1].acceleration
        font_obj = pg.font.Font('freesansbold.ttf', 20)
        text_surface_obj = font_obj.render(
            'vel: {:.2f} km/h '.format(speed), True, (0, 0, 0))
        text_rect_obj = text_surface_obj.get_rect()
        text_rect_obj.center = (110, 180)
        self.screen.blit(text_surface_obj, text_rect_obj)
        text_surface_obj = font_obj.render(
            'acc: {:.2f} m/s^2'.format(acc), True, (0, 0, 0))
        text_rect_obj = text_surface_obj.get_rect()
        text_rect_obj.center = (110, 200)
        self.screen.blit(text_surface_obj, text_rect_obj)

    def plot_ids_on_screen(self):
        for idx in self.recorded_ids:
            state_3dof = (self.vehicles[idx].vec_position.x,
                          self.vehicles[idx].vec_position.y, self.vehicles[idx].angle)
            font_obj = pg.font.Font('freesansbold.ttf', 16)
            text_surface_obj = font_obj.render(
                '{}'.format(idx), True, (0, 255, 0))
            text_rect_obj = text_surface_obj.get_rect()
            u, v = self.project_world_to_image(state_3dof)
            text_rect_obj.center = (u-20, v)
            self.screen.blit(text_surface_obj, text_rect_obj)

    def plot_orientations_on_screen(self):
        for idx in self.recorded_ids:
            state_3dof = (self.vehicles[idx].vec_position.x,
                          self.vehicles[idx].vec_position.y, self.vehicles[idx].angle)
            u, v = self.project_world_to_image(state_3dof)
            angle_diff = self.vehicles[idx].angle - self.center_3dof[2]
            pt1 = (u, v)
            pt2 = (u - 10 * sin(angle_diff), v - 10 * cos(angle_diff))
            pg.draw.line(self.screen, pg.Color('black'), pt1, pt2, 3)

    def plot_selected_rect_on_screen(self):
        if self.agent_id in self.recorded_ids:
            agent_state = (self.vehicles[self.agent_id].vec_position.x,
                           self.vehicles[self.agent_id].vec_position.y, self.vehicles[self.agent_id].angle)
            u, v = self.project_world_to_image(agent_state)
            pg.draw.rect(self.screen, pg.Color('aquamarine3'), (u-10, v-10, 20, 20), 3)

    def init_joy(self, frame_id):
        joy = Joy()
        joy.header.frame_id = frame_id
        joy.header.stamp = self.get_clock().now().to_msg()
        for i in range(8):
            joy.axes.append(0.0)
        for i in range(11):
            joy.buttons.append(0)
        return joy

    def handle_keyboard_event(self):
        for event in pg.event.get():
            if event.type == pg.MOUSEBUTTONUP:
                pos = pg.mouse.get_pos()
                clicked_sprites = [
                    s for s in self.all_sprites if s.rect.collidepoint(pos)]
                if len(clicked_sprites) > 0:
                    if hasattr(clicked_sprites[0], 'id'):
                        self.agent_id = clicked_sprites[0].id
                        self.get_logger().info(f'update agent id to {self.agent_id}')

            if event.type == KEYDOWN:
                joy = self.init_joy("{}".format(self.agent_id))
                if event.key == pg.K_w:
                    msg = 'Agent {}: Speed up'.format(self.agent_id)
                    self.get_logger().info(msg)
                    joy.buttons[3] = 1
                    self.joy_pub.publish(joy)
                elif event.key == pg.K_s:
                    msg = 'Agent {}: Brake'.format(self.agent_id)
                    self.get_logger().info(msg)
                    joy.buttons[0] = 1
                    self.joy_pub.publish(joy)
                elif event.key == pg.K_a:
                    msg = 'Agent {}: Lane change left'.format(self.agent_id)
                    self.get_logger().info(msg)
                    joy.buttons[2] = 1
                    self.joy_pub.publish(joy)
                elif event.key == pg.K_d:
                    msg = 'Agent {}: Lane change right'.format(self.agent_id)
                    self.get_logger().info(msg)
                    joy.buttons[1] = 1
                    self.joy_pub.publish(joy)
                elif event.key == pg.K_q:
                    msg = 'Agent {}: Toggle left lc feasible state'.format(
                        self.agent_id)
                    self.get_logger().info(msg)
                    joy.buttons[4] = 1
                    self.joy_pub.publish(joy)
                elif event.key == pg.K_e:
                    msg = 'Agent {}: Toggle right lc feasible state'.format(
                        self.agent_id)
                    self.get_logger().info(msg)
                    joy.buttons[5] = 1
                    self.joy_pub.publish(joy)
                elif event.key == pg.K_r:
                    msg = 'Agent {}: Toggle autonomous mode'.format(self.agent_id)
                    self.get_logger().info(msg)
                    joy.buttons[6] = 1
                    self.joy_pub.publish(joy)

# Update TerminalServerNode class to inherit from TerminalServerNodeMethods
class TerminalServerNode(Node, TerminalServerNodeMethods):
    def __init__(self):
        super().__init__('terminal_server')
        
        # Global variables moved to class
        self.ego_id = 0
        self.agent_id = 0
        self.vehicles = {}
        self.recorded_ids = []
        self.lane_pts = []
        self.center_3dof = (0.0, 0.0, 0.0)
        self.state_seq = []
        self.has_arena_info_dynamic = False
        
        # Pygame setup
        self.width = 800
        self.height = 600
        self.scale = 4.0
        pg.init()
        self.screen = pg.display.set_mode((self.width, self.height))
        self.all_sprites = pg.sprite.Group()
        
        # ROS2 publishers and subscribers
        self.joy_pub = self.create_publisher(Joy, '/joy', 10)
        
        self.arena_info_dynamic_sub = self.create_subscription(
            ArenaInfoDynamic,
            '/arena_info_dynamic',
            self.process_arena_info_dynamic,
            10
        )
        
        self.arena_info_static_sub = self.create_subscription(
            ArenaInfoStatic,
            '/arena_info_static',
            self.process_arena_info_static,
            10
        )
        
        self.control_signal_sub = self.create_subscription(
            ControlSignal,
            '/ctrl/agent_0',
            self.process_control_signal,
            10
        )
        
        # Timer for main loop
        self.timer = self.create_timer(1.0/50.0, self.main_loop)
        
        # Initialize display
        self.screen.fill(pg.Color('cornsilk3'))
        screen_rect = self.screen.get_rect()
        wheel = Wheel(screen_rect)
        wheel.set_node(self)  # Set reference to node
        self.all_sprites.add(wheel)
        self.all_sprites.draw(self.screen)
        pg.display.set_caption('Ultimate Vehicle Planning')
        pg.display.update()
        
        self.get_logger().info('Terminal server initialized.')


    def project_world_to_image(self, point_3dof):
        x = point_3dof[0] - self.center_3dof[0]
        y = point_3dof[1] - self.center_3dof[1]
        angle = self.center_3dof[2]
        u = self.width/2 + self.scale * (x*sin(angle) - y * cos(angle))
        v = self.height/2 - self.scale * (x*cos(angle) + y * sin(angle))
        return (u, v)

    def process_arena_info_dynamic(self, data):
        for v in data.vehicle_set.vehicles:
            self.vehicles[v.id.data] = v.state
            if v.id.data not in self.recorded_ids:
                self.recorded_ids.append(v.id.data)
                screen_rect = self.screen.get_rect()
                self.all_sprites.add(Vehicle(screen_rect, v.id.data, (v.state.vec_position.x,
                                                                 v.state.vec_position.y, v.state.angle), self))
        self.center_3dof = (self.vehicles[self.ego_id].vec_position.x,
                       self.vehicles[self.ego_id].vec_position.y, self.vehicles[self.ego_id].angle)
        self.has_arena_info_dynamic = True

    def process_arena_info_static(self, data):
        if self.has_arena_info_dynamic:
            visible_range = 150.0
            del self.lane_pts[:]
            for lane in data.lane_net.lanes:
                points = []
                for pt in lane.points:
                    if abs(pt.x - self.center_3dof[0]) < visible_range and abs(pt.y - self.center_3dof[1]) < visible_range:
                        points.append(self.project_world_to_image((pt.x, pt.y, 0.0)))
                if len(points) > 2:
                    self.lane_pts.append(points)

    def process_control_signal(self, data):
        self.state_seq.append(data.state)
        if len(self.state_seq) > 10:
            self.state_seq.pop(0)

    def update_visualization(self):
        if self.has_arena_info_dynamic:
            self.all_sprites.update()
            self.screen.fill(pg.Color('cornsilk2'))
            self.all_sprites.draw(self.screen)
            self.plot_lanes_on_screen()
            self.plot_ids_on_screen()
            self.plot_selected_rect_on_screen()
            self.plot_speed_on_screen()
            self.plot_orientations_on_screen()

    def main_loop(self):
        """Main timer callback for ROS2"""
        self.handle_keyboard_event()
        self.update_visualization()
        pg.display.update()


class Wheel(pg.sprite.Sprite):
    def __init__(self, screen_rect):
        pg.sprite.Sprite.__init__(self)
        # Get the package path for the image
        try:
            package_share = get_package_share_directory('tools')
            image_path = os.path.join(package_share, 'src', 'steer_wheel.png')
        except:
            # Fallback to relative path for development
            image_path = "steer_wheel.png"
        
        self.original_image = pg.image.load(image_path).convert_alpha()
        self.original_image = pg.transform.rotozoom(
            self.original_image, 0.0, 0.2)
        self.image = self.original_image
        self.rect = self.image.get_rect()
        self.rect.center = (100, 100)
        # Store reference to node for calc_current_steer_acc
        self.node = None

    def set_node(self, node):
        self.node = node

    def update(self):
        if self.node:
            angle, acc = self.node.calc_current_steer_acc()
            angle = angle * 180 / pi
            self.image = pg.transform.rotate(self.original_image, angle)
            x, y = self.rect.center
            self.rect = self.image.get_rect()
            self.rect.center = (x, y)


class Vehicle(pg.sprite.Sprite):
    def __init__(self, screen_rect, id, state_3dof, node):
        pg.sprite.Sprite.__init__(self)
        self.id = id
        self.state_3dof = state_3dof
        self.node = node
        self.radius = 10
        self.image = pg.Surface((self.radius*2, self.radius*2), pg.SRCALPHA)
        if id == self.node.ego_id:
            pg.draw.circle(self.image, pg.Color(
                'darkolivegreen1'), (self.radius, self.radius), self.radius)
        else:
            pg.draw.circle(self.image, pg.Color(
                'dodgerblue1'), (self.radius, self.radius), self.radius)

        self.rect = self.image.get_rect(
            center=self.node.project_world_to_image(state_3dof))
        self.screen_rect = screen_rect

    def update(self):
        latest_state_3dof = (self.node.vehicles[self.id].vec_position.x,
                             self.node.vehicles[self.id].vec_position.y, self.node.vehicles[self.id].angle)
        self.rect.center = self.node.project_world_to_image(latest_state_3dof)






def main():
    rclpy.init()
    node = TerminalServerNode()
    
    try:
        rclpy.spin(node)
    except KeyboardInterrupt:
        pass
    finally:
        node.destroy_node()
        rclpy.shutdown()
        pg.quit()


if __name__ == '__main__':
    main()
